/*
 * To change this template, choose Tools | Templates
 * and open the template in the editor.
 */
package technotheist.imgcog;

import com.amd.aparapi.Kernel;
import java.util.ArrayList;

/**
 *
 * @author Mitch
 */
public abstract class Net implements DataProvider, DataListener {

    private DataListener dataListener;
    private boolean lastDirty = true;
    private NetRunner runner;

    @Override
    public boolean isCompatible(DataProvider dp) {
        return dp.getOutputLength() == runner.firstLayerStore.length;
    }

    @Override
    public boolean addListener(DataListener listener) {
        if (dataListener == null && listener.isCompatible(this)) {
            dataListener = listener;
            return true;
        }
        return false;
    }

    @Override
    public boolean removeListener(DataListener listener) {
        if(dataListener == listener) {
            dataListener = null;
            return true;
        }
        return false;
    }

    @Override
    public void dataUpdateEvent(DataProvider dp) {
        if (setData(dp)) {
            forward();
            dataListener.dataUpdateEvent(this);
        }
    }

    @Override
    public DimType getOutputDimType() {
        return DimType.Vec1D;
    }

    @Override
    public int getOutputDimLen(int i) {
        if (i == 0) {
            return runner.layerSize[0];
        }
        return -1;
    }

    @Override
    public int getOutputLength() {
        return runner.layerSize[runner.layerCount - 1];
    }

    protected boolean setData(DataProvider dp) {
        float[] src = dp.getData();
        if (src.length == runner.firstLayerStore.length) {
            System.arraycopy(src, 0, runner.firstLayerStore, 0, src.length);
            runner.put(runner.firstLayerStore);
            return true;
        }
        return false;
    }

    @Override
    public float[] getData() {
        if (lastDirty) {
            lastDirty = false;
            runner.get(runner.lastLayerStore);
        }
        return runner.lastLayerStore;
    }

    protected NetRunner getRunner() {
        return runner;
    }

    protected void forward() {
        lastDirty = true;
    }

    public interface NetConfig {

        public int getLayerCount();

        public int getLayerSize(int layer);

        public int getWeightsPerNode(int layer);

        public int getTotalWeightCount();

        public int getTotalNodeCount();

        public int getNumWeights(int i);

        public boolean sharedBias();
    }

    //
    protected abstract class NetRunner extends Kernel {

        protected final int layerCount;
        @Constant
        private final int[] layerSize;
        @Constant
        private final int[] hiddenLayerStart;
        //
        private final float[] firstLayerStore;
        private final float[] lastLayerStore;
        private final float[] hiddenLayerStore;
        //
        @Constant
        private final int[] weightStart;

        protected NetRunner(NetRunner runner) {
            this.layerCount = runner.layerCount;
            this.layerSize = runner.layerSize;
            this.hiddenLayerStart = runner.hiddenLayerStart;
            this.firstLayerStore = runner.firstLayerStore;
            this.lastLayerStore = runner.lastLayerStore;
            this.hiddenLayerStore = runner.hiddenLayerStore;
            this.weightStart = runner.weightStart;
        }

        protected NetRunner(NetConfig conf) {
            layerCount = conf.getLayerCount();
            int sum = 0;
            int sumH = 0;
            int sumW = 0;
            layerSize = new int[layerCount];
            hiddenLayerStart = new int[layerCount - 1];
            for (int i = 0; i < layerCount; i++) {
                layerSize[i] = conf.getLayerSize(i);
                sum += layerSize[i];
                if (i > 0) {
                    sumW += conf.getWeightsPerNode(i);
                    if (i < layerCount - 1) {
                        hiddenLayerStart[i - 1] = sumH;
                        sumH += layerSize[i];
                    }
                }
            }
            firstLayerStore = new float[layerSize[0]];
            lastLayerStore = new float[layerSize[layerCount - 1]];
            hiddenLayerStore = new float[sumH];
            weightStart = new int[layerCount - 1];
            for (int i = 0; i < layerCount - 1; i++) {
                weightStart[i] = conf.getNumWeights(i);
            }
        }

        protected int getWeightStart(int layerA) {
            return weightStart[layerA];
        }

        protected int getHiddenStart(int layer) {
            return hiddenLayerStart[layer - 1];
        }

        protected final int getLayerSize(int layer) {
            return layerSize[layer];
        }

        protected final float getNodeStore(int layer, int node) {
            if (layer == 0) {
                return firstLayerStore[node];
            }
            if (layer == layerCount - 1) {
                return lastLayerStore[node];
            }
            if (layer < layerCount - 1) {
                return hiddenLayerStore[hiddenLayerStart[layer - 1] + node];
            }
            return 0;
        }

        protected final void setNodeStore(int layer, int node, float val) {
            if (layer == 0) {
                firstLayerStore[node] = val;
            }
            if (layer == layerCount - 1) {
                lastLayerStore[node] = val;
            }
            if (layer < layerCount - 1) {
                hiddenLayerStore[hiddenLayerStart[layer - 1] + node] = val;
            }
        }

        protected final void addNodeStore(int layer, int node, float val) {
            if (layer == 0) {
                firstLayerStore[node] = +val;
            }
            if (layer == layerCount - 1) {
                lastLayerStore[node] = +val;
            }
            if (layer < layerCount - 1) {
                hiddenLayerStore[hiddenLayerStart[layer - 1] + node] = +val;
            }
        }

        protected abstract int getInputWeightIndex(int layer, int node, int w);

        protected abstract float getWeight(int index);
    }

    protected abstract class TrainedNetRunner extends NetRunner {

        @Constant
        private final float[] weights;
        @Constant
        private final float[] bias;

        public TrainedNetRunner(UntrainedNetRunner runner) {
            super(runner);
            runner.get(runner.weights);
            weights = runner.weights.clone();
            bias = runner.bias.clone();
        }

        @Override
        protected float getWeight(int index) {
            return weights[index];
        }
    }

    protected abstract class UntrainedNetRunner extends NetRunner {

        private final float[] weights;
        private final float[] bias;

        public UntrainedNetRunner(NetConfig conf) {
            super(conf);
            weights = new float[conf.getTotalWeightCount()];
            if (conf.sharedBias()) {
                bias = new float[layerCount - 1];
            } else {
                bias = new float[conf.getTotalNodeCount() - getLayerSize(0)];
            }
        }

        @Override
        protected float getWeight(int index) {
            return weights[index];
        }

        protected void setWeight(int index, float val) {
            weights[index] = val;
        }

        protected void addWeight(int index, float val) {
            weights[index] += val;
        }
    }
}